#ifndef DECISION_STUMP_H_
#define DECISION_STUMP_H_

#include "WeakLearner.h"
#include <set>

GLOBAL_NAMESPACE_BEGIN

NAMESPACE_MACHINE_LEARNING_BEGIN

class DecisionStump : public WeakLearner
{
public:
    struct Parameters
    {
        std::set<int> subClass;
        int           dimF;
        bool          bGreater;
        double        threshold;
        double        a;
        double        b;
        double        k;
    };
public:
    DecisionStump();
    DecisionStump(const DecisionStump& rhs);
    ~DecisionStump();

    void reset();
    void setSubClasses(const std::set<int>& classes, int maxLabel, int thisClassIdx);
    void setMaxMinCoeffs(const std::vector<double>& maxCoeffs, const std::vector<double>& minCoeffs, int stepNum);
    void setWeightPrecomputeData(const std::vector<double>& t0c, double totalWeight);

    int     train    (const Eigen::MatrixXd& trainData, const Eigen::VectorXi& labels, const Eigen::MatrixXd& weights);
    int     predict  (const Eigen::RowVectorXd& oneData) const;
    double  evaluate (const Eigen::RowVectorXd& oneData, int classIdx) const;
    double  evaluate (const Eigen::MatrixXd& data, int dataIdx, int classIdx) const;

    //int     train   (const Eigen::MatrixXd& trainData, const Eigen::MatrixXi& labels, const Eigen::MatrixXd& weights);
    //Eigen::RowVectorXi predict(const Eigen::RowVectorXd& oneData) const;

    double  evalError(const Eigen::MatrixXd& data, const Eigen::VectorXi& labels, const Eigen::MatrixXd& weights) const;
    int  getLabelResponse(int label) const;

    void output(std::ostream& out);

    const Parameters* getParameters() const;

private:
    //std::set<int> subClass_;
    //int         dimF_;
    //bool        bGreater_;
    //double      threshold_;
    //double      a_;
    //double      b_;
    //double      k_;
    Parameters param_;
    Parameters tmpParam_;   // this is for training

    int         maxLabel_;
    int         classIdx_;
    int         stepNum_;
    const std::vector<double>* maxCoeffs_;
    const std::vector<double>* minCoeffs_;
    const std::vector<double>* t0c_;    // total weights of class c
    double totalWeight_;                // total weights of all classes

private:
    bool isLabelInStump(int label, bool bUseTmpParams = false) const;
    int  getLabelResponse(int label, bool bUseTmpParameter);

    int buildStump(const Eigen::MatrixXd& trainData, const Eigen::VectorXi& labels, const Eigen::MatrixXd& weights,
        const std::vector<double>& maxCoeffs, const std::vector<double>& minCoeffs, int stepNum);
    int trainStump(const Eigen::MatrixXd& trainData, const Eigen::VectorXi& labels, const Eigen::MatrixXd& weights);

    int buildStump(const Eigen::MatrixXd& trainData, const Eigen::VectorXi& labels, const Eigen::MatrixXd& weights,
        const std::vector<double>& maxCoeffs, const std::vector<double>& minCoeffs, int stepNum, bool method2);
    int trainStump(const Eigen::MatrixXd& trainData, const Eigen::VectorXi& labels, const Eigen::MatrixXd& weights, bool method2);
    double evalError(const Eigen::MatrixXd& data, const Eigen::VectorXi& labels, const Eigen::MatrixXd& weights, bool method2);
    double evaluate (const Eigen::MatrixXd& data, int dataIdx, int classIdx, bool method2);

};

NAMESPACE_MACHINE_LEARNING_END

GLOBAL_NAMESPACE_END

#endif//DECISION_STUMP_H_